#include "code_comment_visitor.h"

bool FuncCommentGenVisitor::TraverseCXXMethodDecl(clang::CXXMethodDecl *D)
{
    if(!D->hasBody()) {
        return true;
    }
    if("__fakeHelperFuncForCommentGen__" == D->getQualifiedNameAsString()) {
        return true;
    }

    mCurrentCommentInfoPtr = getCommentInfo(D);
    CodeCommentInfoNode* newNode = createNewNodeWithType(METHOD_FUNC_COMMENT);
    if (nullptr == newNode) {
        return true;
    }

    newNode->mKeyInfo["Name"] = D->getQualifiedNameAsString();
    clang::CXXRecordDecl* cxxDecl = D->getParent();
    if (cxxDecl) {
        newNode->mKeyInfo["Class"] = cxxDecl->getQualifiedNameAsString();
    }
    int i = 0;
    for(auto paraDecl : D->parameters()) {
        newNode->mKeyInfo[std::string("p") + std::to_string(i)] = paraDecl->getQualifiedNameAsString();
        i++;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    TraverseStmt(D->getBody());
    mCurrentCommentInfoPtr->pop();
    mCurrentCommentInfoPtr = nullptr;
    return true;
}

bool FuncCommentGenVisitor::TraverseFunctionDecl(clang::FunctionDecl *D)
{   
    if(!D->hasBody()) {
        return true;
    }

    if("__fakeHelperFuncForCommentGen__" == D->getQualifiedNameAsString()) {
        return true;
    }

    mCurrentCommentInfoPtr = getCommentInfo(D);
    CodeCommentInfoNode* newNode = createNewNodeWithType(FUNC_COMMENT);
    if (nullptr == newNode) {
        return true;
    }
    
    newNode->mKeyInfo["Name"] = D->getQualifiedNameAsString();
    int i = 0;
    for(auto paraDecl : D->parameters()) {
        newNode->mKeyInfo[std::string("p") + std::to_string(i)] = paraDecl->getQualifiedNameAsString();
        i++;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    TraverseStmt(D->getBody());
    mCurrentCommentInfoPtr->pop();
    mCurrentCommentInfoPtr = nullptr;
     
    return true;
}
    
bool FuncCommentGenVisitor::TraverseIfStmt(clang::IfStmt *stmt)
{   
    CodeCommentInfoNode* newNode = createNewNodeWithType(IF_THEN_COMMENT);
    if (nullptr == newNode) {
        return true;
    }
    
    clang::Expr* cond = stmt->getCond();
    if(cond) {
        std::string thenStr;
        getStmtStr(cond, thenStr);
        newNode->mKeyInfo["then"] = thenStr;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* thenStmt = stmt->getThen();
    if (thenStmt) {
        TraverseStmt(thenStmt);
    }
    mCurrentCommentInfoPtr->pop();
    
    clang::Stmt* elseStmt = stmt->getElse();
    if(elseStmt) {
        CodeCommentInfoNode* newNode = createNewNodeWithType(ELSE_COMMENT);
        if (nullptr == newNode) {
            return true;
        }
        mCurrentCommentInfoPtr->pushInfo(newNode);
        TraverseStmt(elseStmt);
        mCurrentCommentInfoPtr->pop();
    }
    return true;
}

bool FuncCommentGenVisitor::TraverseForStmt(clang::ForStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(FOR_COMMENT);
    if (nullptr == newNode) {
        return true;
    }
    
    const clang::Stmt* initStmt = stmt->getInit();
    const clang::Expr* condExpr = stmt->getCond();
    const clang::Expr* incExpr = stmt->getInc();
    if(initStmt) {
        std::string initStr;
        getStmtStr(initStmt, initStr);
        newNode->mKeyInfo["init"] = initStr;
    }
    if(condExpr) {
        std::string condStr;
        getStmtStr(condExpr, condStr);
        newNode->mKeyInfo["cond"] = condStr;
    }
    if(incExpr) {
        std::string incStr;
        getStmtStr(incExpr, incStr);
        newNode->mKeyInfo["inc"] = incStr;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* bodyStmt = stmt->getBody();
    if(bodyStmt) {
        TraverseStmt(bodyStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    
    return true;
}

bool FuncCommentGenVisitor::TraverseWhileStmt(clang::WhileStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(WHILE_COMMENT);
    if (nullptr == newNode) {
        return true;
    }
    
    clang::Expr* condExpr = stmt->getCond();
    if(condExpr) {
        std::string condStr;
        getStmtStr(condExpr, condStr);
        newNode->mKeyInfo["cond"] = condStr;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* bodyStmt = stmt->getBody();
    if(bodyStmt) {
        TraverseStmt(bodyStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    
    return true;
}

bool FuncCommentGenVisitor::TraverseDoStmt(clang::DoStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(DO_WHILE_COMMENT);
    if (nullptr == newNode) {
        return true;
    }

    clang::Expr* condExpr = stmt->getCond();
    if(condExpr) {
        std::string condStr;
        getStmtStr(condExpr, condStr);
        newNode->mKeyInfo["cond"] = condStr;
    }
    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* bodyStmt = stmt->getBody();
    if(bodyStmt) {
        TraverseStmt(bodyStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    
    return true;
}

bool FuncCommentGenVisitor::TraverseSwitchStmt(clang::SwitchStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(SWITCH_COMMENT);
    if (nullptr == newNode) {
        return true;
    }

    // llvm::outs() << "SWITCH STMT .\n";
    // stmt->dump();

    clang::Expr* condExpr = stmt->getCond();
    if(condExpr) {
        std::string condStr;
        getStmtStr(condExpr, condStr);
        newNode->mKeyInfo["cond"] = condStr;
    }
     
    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* bodyStmt = stmt->getBody();
    if(bodyStmt) {
        TraverseStmt(bodyStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    
    return true;
}

bool FuncCommentGenVisitor::TraverseCaseStmt(clang::CaseStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(CASE_COMMENT);
    if (nullptr == newNode) {
        return true;
    }
    
    clang::Expr* lhsExpr = stmt->getLHS();
    if(lhsExpr) {
        std::string lhsStr;
        getStmtStr(lhsExpr, lhsStr);
        newNode->mKeyInfo["cond"] = lhsStr;
        // llvm::outs() << "LHS. \n";
        // lhsExpr->dump();
    }
    clang::Expr* rhsExpr = stmt->getRHS();
    if(rhsExpr) {
        std::string rhsStr;
        getStmtStr(rhsExpr, rhsStr);
        newNode->mKeyInfo["cond_right"] = rhsStr;
        // llvm::outs() << "RHS. \n";
        // rhsExpr->dump();
    }

    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* subStmt = stmt->getSubStmt();
    if(subStmt) {
        // llvm::outs() << "CASE SUB. \n";
        // subStmt->dump();
        TraverseStmt(subStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    return true;
}

bool FuncCommentGenVisitor::TraverseDefaultStmt(clang::DefaultStmt *stmt)
{
    CodeCommentInfoNode* newNode = createNewNodeWithType(DEFAULT_COMMENT);
    if (nullptr == newNode) {
        return true;
    }

    mCurrentCommentInfoPtr->pushInfo(newNode);
    clang::Stmt* subStmt = stmt->getSubStmt();
    if(subStmt) {
        TraverseStmt(subStmt);   
    }
    mCurrentCommentInfoPtr->pop();
    return true;
}


bool FuncCommentGenVisitor::VisitCallExpr(const clang::CallExpr *expr)
{
    const clang::FunctionDecl* funcDecl = expr->getDirectCallee();
    if (nullptr == funcDecl || 2 > expr->getNumArgs()) {
        return true;
    }
    
    if("__fakeHelperFuncForCommentGen__" != funcDecl->getQualifiedNameAsString()) {
        return true;
    }

    const clang::Expr* argExpr = expr->getArg(1);
    if(argExpr) {
        std::string argStr;
        getStmtStr(argExpr, argStr);
        CodeCommentInfoNode* newNode = createNewNodeWithType(ACTION_COMMENT);
        if (nullptr == newNode) {
            return true;
        }
        newNode->mKeyInfo["action"] = argStr;
        mCurrentCommentInfoPtr->addInfo(newNode);
    }
    return true;
}

bool FuncCommentGenVisitor::getStmtStr(const clang::Stmt* stmt, std::string& stmtStr)
{
    llvm::raw_string_ostream OS(stmtStr);
    stmt->printPretty(OS, nullptr, mContext.getPrintingPolicy());
    return true;
}

bool FuncCommentGenVisitor::getFile(const clang::NamedDecl *D, const clang::ASTContext &Context,
                                            std::string& fileName)
{
    fileName = std::string(Context.getSourceManager()
                                  .getPresumedLoc(D->getBeginLoc())
                                  .getFilename());
    return true;
}

int FuncCommentGenVisitor::getLine(const clang::NamedDecl *D, const clang::ASTContext &Context)
{
    return Context.getSourceManager().getPresumedLoc(D->getBeginLoc()).getLine();
}

CodeCommentInfo* FuncCommentGenVisitor::getCommentInfo(const clang::NamedDecl* D)
{
    std::string fileIdent;
    getFile(D, D->getASTContext(), fileIdent);
    std::string funcKey = D->getQualifiedNameAsString();
    return mCodeCommentInfoManager.getOrAddCodeCommentInfo(fileIdent, funcKey);
}

CodeCommentInfoNode* FuncCommentGenVisitor::createNewNodeWithType(const CommentType ct)
{
    if(nullptr == mCurrentCommentInfoPtr) {
        // if we can not locate current node, there maybe an error occurs
        llvm::errs() << "ERROR: can not locate current CodeCommentInfo, Please Check. \n";
        return nullptr;
    }
    // set new CodeCommentInfoNode
    CodeCommentInfoNode* newNode = new CodeCommentInfoNode();
    if(nullptr == newNode) {
        return newNode;
    }
    
    newNode->mCommentType = ct;
    return newNode;
}
